#include <iostream>
#include <graphics.h>
#include <vector>
#include <math.h>

using namespace std;

int width = 1024, height = 786;
int bricks_count = 0;
void koef();

int colourfunc( float a )
{
  if(a == 1) return 15;
  if(a < 1 && a >= 0.95) return 4;
  if(a < 0.95 && a >= 0.90) return 6;
  if(a < 0.90 && a >= 0.85) return 2;
  if(a < 0.85 && a >= 0.01) return 3;
  return 0;
}

void moveto1( float x, float y){moveto((int)x, (int)y);}

void lineto1( float x, float y){lineto((int)x, (int)y);}

class vec2
{
public:      
  
  float x, y;

  vec2( float nx = 0, float ny  = 0):
         x(nx), y(ny)
  {}
  
  vec2 operator +( vec2 & v )
  {
    return vec2(x + v.x, y + v.y);
  }

  vec2 operator -( vec2 & v )
  {
    return vec2(x - v.x, y - v.y);
  }

  vec2 operator *( float n )
  {
    return vec2(x * n, y * n);
  }

  vec2 operator *( vec2 &v )
  {
    return vec2(x * v.x + y * v.y);
  }
  
  float length()
  {
    return sqrt(x * x + y * y);
  }
};

class edge
{
public:     
  float y, n1, n2;

  edge( float ny, float nn1 = 1, float nn2 = 1 ) : 
    y(ny), n1(nn1), n2(nn2)
  {
    moveto1(0, y);
    lineto1(width, y);
  }
};

vector<edge> edges;

void surface( int bricks_count )
{  	 	 
  for(int i = 0; i <= bricks_count; ++i)
  {
    edge new_edge(height * (1 / 3.0f + i / (bricks_count * 3.0f)));
    edges.push_back(new_edge);
  } 
  koef();       
}

void koef()
{
  for(int i = 0; i <= bricks_count; ++i)
  {
    edges[i].n1 = 1 + sin(3.141592653589793 * i / 6.22) * sin(3.141592653589793 * i / 6.22);  
    edges[i].n2 = 1 + sin(3.141592653589793 * (i + 1) / 6.22) * sin(3.141592653589793 * (i + 1) / 6.22);  
  }
  edges[bricks_count].n2 = 1;  
}
  
void put_ray( vec2 pos, vec2 dir, int edge_index, float intensity )
{
  if(intensity < 0.01 || pos.x >= width)
    return;
  
  //fall
  setcolor(colourfunc(intensity));
  moveto1(pos.x, pos.y); 
  lineto1(dir.x, dir.y);	
  
  if (edge_index < 0 || edge_index > bricks_count)
    return;
  
  float n1 = edges[edge_index].n1, n2 = edges[edge_index].n2, temp;
  int i = 1;//norm
  vec2 ray(dir - pos);
  
  if(ray.y < 0)
  {
    temp = n1;
    n1 = n2;
    n2 = temp;
    i = -1;
  }   

  //full reflection
  if(ray.x / ray.length() * n1 >= n2)
  {
    put_ray(dir, vec2(pos.x + ray.x * 2.0f, pos.y), edge_index - i, intensity);
    return;
  }  

  //reflection intensity
  float ref_in = n1 / n2;
  float ref_in2 = ref_in * ref_in;
  float cos_a = sqrt(ray.y * ray.y) / ray.length()
       ,c     = cos_a * ref_in;
  float g = sqrt(1.0f + c * c - ref_in2);  
  float rp = (n1 * g - n2 * cos_a) / (n1 * g + n2 * cos_a) * (n1 * g - n2 * cos_a) / (n1 * g + n2 * cos_a); // coefficient reflection energy
  float rs = (n2 * g - n1 * cos_a) / (n2 * g + n1 * cos_a) * (n2 * g - n1 * cos_a) / (n2 * g + n1 * cos_a);
  float ra = 0.5f * (rp + rs);
  
  cout << "reflection : " << ra << endl;
  cout << "refraction : " << (1.0f - ra) << endl;
  
  //reflection
  put_ray(dir, vec2(pos.x + ray.x * 2.0f, pos.y), edge_index - i, intensity * ra);
  
  //refraction
  vec2 ray1;
  ray1.y = (height /(bricks_count * 3.0f));
  if((edge_index == 0 && i == -1) || (edge_index == bricks_count && i == 1))
    ray1.y = height / 3.0f;
    
  float sin_b = n1 / n2 * ray.x / ray.length();
  float tan_b = sin_b / sqrt(1.0f - sin_b * sin_b);
  
  ray1.x = ray1.y * tan_b;
  ray1.y = ray1.y * i;
  
  put_ray(dir, dir + ray1, edge_index + i, intensity * (1.0f - ra));
}

int main()
{   
  cout << "Input number of bricks: ";
  cin >> bricks_count;

  initwindow(width, height);

  surface(bricks_count);
  put_ray(vec2(0.0f,0.0f), vec2(width / 5.0f, height / 3.0f), 0, 1.0f);

  getch();
  closegraph();
}
